package timing.ukulele.auth.authorization;

import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.source.ImmutableJWKSet;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;
import lombok.SneakyThrows;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cglib.beans.BeanMap;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;
import org.springframework.http.MediaType;
import org.springframework.jdbc.core.JdbcOperations;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.config.annotation.authentication.configuration.AuthenticationConfiguration;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.http.SessionCreationPolicy;
import org.springframework.security.oauth2.core.oidc.OidcUserInfo;
import org.springframework.security.oauth2.jwt.JwtDecoder;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.NimbusJwtEncoder;
import org.springframework.security.oauth2.server.authorization.*;
import org.springframework.security.oauth2.server.authorization.client.JdbcRegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.config.annotation.web.configuration.OAuth2AuthorizationServerConfiguration;
import org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2AuthorizationServerConfigurer;
import org.springframework.security.oauth2.server.authorization.oidc.authentication.OidcUserInfoAuthenticationContext;
import org.springframework.security.oauth2.server.authorization.oidc.authentication.OidcUserInfoAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.settings.AuthorizationServerSettings;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationConverter;
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationToken;
import org.springframework.security.oauth2.server.resource.authentication.JwtGrantedAuthoritiesConverter;
import org.springframework.security.rsa.crypto.KeyStoreKeyFactory;
import org.springframework.security.web.SecurityFilterChain;
import org.springframework.security.web.header.HeaderWriterFilter;
import org.springframework.security.web.util.matcher.MediaTypeRequestMatcher;
import org.springframework.util.ObjectUtils;
import timing.ukulele.auth.authorization.filter.JWTAuthorizationFilter;
import timing.ukulele.auth.authorization.handler.ConsentAuthenticationFailureHandler;
import timing.ukulele.auth.authorization.handler.ConsentAuthorizationResponseHandler;
import timing.ukulele.auth.authorization.handler.LoginTargetAuthenticationEntryPoint;
import timing.ukulele.auth.authorization.token.TimingJwtGenerator;
import timing.ukulele.auth.config.properties.TimingSecurityProperties;
import timing.ukulele.auth.constant.RedisConstants;
import timing.ukulele.auth.security.user.UnionUserDetailService;
import timing.ukulele.auth.support.RedisOperator;
import timing.ukulele.common.data.ResponseData;
import timing.ukulele.data.user.view.UserVO;

import java.security.KeyPair;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.function.Function;

import static timing.ukulele.share.Constant.AUTHORITIES_KEY;

@Configuration
public class AuthorizationConfig {

    private final RedisOperator<String> redisOperator;
    private final TimingSecurityProperties securityProperties;
    private final ResourceLoader resourceLoader;
    private final UnionUserDetailService userDetailService;
    @Value("${timing.auth.jksPath}")
    private String jksPath;
    @Value("${timing.auth.password}")
    private String jksPassword;
    @Value("${timing.auth.alias}")
    private String jksAlias;

    public AuthorizationConfig(
            RedisOperator<String> redisOperator,
            TimingSecurityProperties securityProperties,
            ResourceLoader resourceLoader,
            UnionUserDetailService userDetailService) {
        this.redisOperator = redisOperator;
        this.securityProperties = securityProperties;
        this.resourceLoader = resourceLoader;
        this.userDetailService = userDetailService;
    }

    @Bean
    public TimingJwtGenerator jwtGenerator(JwtEncoder jwtEncoder) {
        TimingJwtGenerator jwtGenerator = new TimingJwtGenerator(jwtEncoder, redisOperator);
//        jwtGenerator.setJwtCustomizer(context -> {
//            OAuth2Authorization authorization = context.getAuthorization();
//            String principalName = authorization.getPrincipalName();
//            System.out.println(principalName);
//        });
        return jwtGenerator;
    }

    /**
     * 自定义jwt解析器，设置解析出来的权限信息的前缀与在jwt中的key
     *
     * @return jwt解析器 JwtAuthenticationConverter
     */
    @Bean
    public JwtAuthenticationConverter jwtAuthenticationConverter() {
        JwtGrantedAuthoritiesConverter grantedAuthoritiesConverter = new JwtGrantedAuthoritiesConverter();
        // 设置解析权限信息的前缀，设置为空是去掉前缀
        grantedAuthoritiesConverter.setAuthorityPrefix("");
        // 设置权限信息在jwt claims中的key
        grantedAuthoritiesConverter.setAuthoritiesClaimName(AUTHORITIES_KEY);

        JwtAuthenticationConverter jwtAuthenticationConverter = new JwtAuthenticationConverter();
        jwtAuthenticationConverter.setJwtGrantedAuthoritiesConverter(grantedAuthoritiesConverter);

        return jwtAuthenticationConverter;
    }


    /**
     * 将AuthenticationManager注入ioc中，其它需要使用地方可以直接从ioc中获取
     *
     * @param authenticationConfiguration 导出认证配置
     * @return AuthenticationManager 认证管理器
     */
    @Bean
    @SneakyThrows
    public AuthenticationManager authenticationManager(AuthenticationConfiguration authenticationConfiguration) {
        return authenticationConfiguration.getAuthenticationManager();
    }

    /**
     * 配置jwk源，使用非对称加密，公开用于检索匹配指定选择器的JWK的方法
     *
     * @return JWKSource
     */
    @Bean
    @SneakyThrows
    public JWKSource<SecurityContext> jwkSource() {
        // 先从redis获取
        String jwkSetCache = redisOperator.get(RedisConstants.AUTHORIZATION_JWS_PREFIX_KEY);
        if (ObjectUtils.isEmpty(jwkSetCache)) {
            KeyPair keyPair = generateRsaKey();
            RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
            RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
            RSAKey rsaKey = new RSAKey.Builder(publicKey)
                    .privateKey(privateKey)
                    .keyID(UUID.randomUUID().toString())
                    .build();
            // 生成jws
            JWKSet jwkSet = new JWKSet(rsaKey);
            // 转为json字符串
            String jwkSetString = jwkSet.toString(Boolean.FALSE);
            // 存入redis
            redisOperator.set(RedisConstants.AUTHORIZATION_JWS_PREFIX_KEY, jwkSetString);
            return new ImmutableJWKSet<>(jwkSet);
        }
        // 解析存储的jws
        JWKSet jwkSet = JWKSet.parse(jwkSetCache);
        return new ImmutableJWKSet<>(jwkSet);
    }

    /**
     * 生成rsa密钥对，提供给jwk
     *
     * @return 密钥对
     */
    private KeyPair generateRsaKey() {
        KeyPair keyPair;
        try {
            Resource resource = resourceLoader.getResource(jksPath);
            KeyStoreKeyFactory factory = new KeyStoreKeyFactory(
                    resource, jksPassword.toCharArray());
            keyPair = factory.getKeyPair(jksAlias);
        } catch (Exception ex) {
            throw new IllegalStateException(ex);
        }
        return keyPair;
    }

    /**
     * 配置jwt解析器
     *
     * @param jwkSource jwk源
     * @return JwtDecoder
     */
    @Bean
    public JwtDecoder jwtDecoder(JWKSource<SecurityContext> jwkSource) {
        return OAuth2AuthorizationServerConfiguration.jwtDecoder(jwkSource);
    }

    @Bean
    JwtEncoder jwtEncoder(JWKSource<SecurityContext> jwkSource) {
        return new NimbusJwtEncoder(jwkSource);
    }

    /**
     * 添加认证服务器配置，设置jwt签发者、默认端点请求地址等
     *
     * @return AuthorizationServerSettings
     */
    @Bean
    public AuthorizationServerSettings authorizationServerSettings() {
        // 注意！！！！此处的Issuer很重要，资源服务里的配置要与此一致
        return AuthorizationServerSettings.builder().issuer(securityProperties.getIssuerUrl()).build();
    }

    /**
     * 配置端点的过滤器链
     *
     * @param http spring security核心配置类
     * @return 过滤器链
     * @throws Exception 抛出
     */
    @Bean
    public SecurityFilterChain authorizationServerSecurityFilterChain(HttpSecurity http,
                                                                      JwtDecoder jwtDecoder) throws Exception {
        // 配置默认的设置，忽略认证端点的csrf校验
        OAuth2AuthorizationServerConfiguration.applyDefaultSecurity(http);
        http.sessionManagement(session -> session.sessionCreationPolicy(SessionCreationPolicy.STATELESS));
        //自定义用户映射器
        Function<OidcUserInfoAuthenticationContext, OidcUserInfo> userInfoMapper = (context) -> {
            OidcUserInfoAuthenticationToken authentication = context.getAuthentication();
            JwtAuthenticationToken principal = (JwtAuthenticationToken) authentication.getPrincipal();
            ResponseData<UserVO> userRes = userDetailService.getUserService().getUserByUserName(principal.getName());
            if (userRes == null || userRes.getData() == null)
                return null;
            BeanMap beanMap = BeanMap.create(userRes.getData());
            Map<String, Object> userMap = new HashMap<>(beanMap.size());
            beanMap.putAll(userMap);
            return new OidcUserInfo(userMap);
        };
        ConsentAuthorizationResponseHandler consentAuthorizationResponseHandler = new ConsentAuthorizationResponseHandler(securityProperties.getConsentPageUri());
        http.getConfigurer(OAuth2AuthorizationServerConfigurer.class)
                // 开启OpenID Connect 1.0协议相关端点
//                .oidc(Customizer.withDefaults())
                .oidc((oidc) -> oidc.userInfoEndpoint((userInfo) -> {
                            userInfo.userInfoMapper(userInfoMapper);
                            userInfo.userInfoResponseHandler(consentAuthorizationResponseHandler);
                        }
                ))
                // 设置自定义用户确认授权页
                .authorizationEndpoint(authorizationEndpoint -> {
//                            // 如果是分离页面则重定向，否则转发请求
                            authorizationEndpoint.consentPage(securityProperties.getConsentPageUri());
                            // 适配前后端分离的授权确认页面，成功/失败响应json
                            authorizationEndpoint.errorResponseHandler(new ConsentAuthenticationFailureHandler(securityProperties.getConsentPageUri()));
                            authorizationEndpoint.authorizationResponseHandler(consentAuthorizationResponseHandler);
                        }
                );

        http
                // 当未登录时访问认证端点时重定向至login页面
                .exceptionHandling((exceptions) -> exceptions
                        .defaultAuthenticationEntryPointFor(
                                new LoginTargetAuthenticationEntryPoint(securityProperties.getLoginUrl(), securityProperties.getDeviceActivateUri()),
                                new MediaTypeRequestMatcher(MediaType.TEXT_HTML)
                        )
                );
        JWTAuthorizationFilter jwtFilter = new JWTAuthorizationFilter(redisOperator, jwtDecoder);
//        http.addFilterAfter(jwtFilter, SecurityContextHolderFilter.class);
        http.addFilterBefore(jwtFilter, HeaderWriterFilter.class);
        return http.build();
    }

    /**
     * 配置客户端Repository
     *
     * @return 基于数据库的repository
     */
    @Bean
    public RegisteredClientRepository registeredClientRepository(JdbcOperations jdbcOperations) {
        // 基于db存储客户端，还有一个基于内存的实现 JdbcRegisteredClientRepository
//        JdbcRegisteredClientRepository registeredClientRepository = new JdbcRegisteredClientRepository(jdbcOperations);
        return new JdbcRegisteredClientRepository(jdbcOperations);
    }

    /**
     * oauth2的授权管理服务
     *
     * @param registeredClientRepository 上边注入的客户端repository
     * @return JdbcOAuth2AuthorizationService
     */
    @Bean
    public OAuth2AuthorizationService authorizationService(RegisteredClientRepository registeredClientRepository,
                                                           JdbcOperations jdbcOperations) {
        // 基于db的oauth2认证服务，JdbcOAuth2AuthorizationService InMemoryOAuth2AuthorizationService
        return new JdbcOAuth2AuthorizationService(jdbcOperations, registeredClientRepository);
    }

    /**
     * 配置基于db的授权确认管理服务
     *
     * @param registeredClientRepository 客户端repository
     * @return JdbcOAuth2AuthorizationConsentService
     */
    @Bean
    public OAuth2AuthorizationConsentService authorizationConsentService(RegisteredClientRepository registeredClientRepository,
                                                                         JdbcOperations jdbcOperations) {
        // 基于db的授权确认管理服务，还有一个基于内存的服务实现：InMemoryOAuth2AuthorizationConsentService JdbcOAuth2AuthorizationConsentService
        return new JdbcOAuth2AuthorizationConsentService(jdbcOperations, registeredClientRepository);
    }

}
